//
// Created by xu on 2024/11/21.
//
#include <stdio.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/imgutils.h>

// 错误处理
#define ERROR_BUF(ret) \
    char errbuf[1024]; \
    av_strerror(ret, errbuf, sizeof (errbuf));


typedef struct {
    const char *filename;
    int width;
    int height;
    enum AVPixelFormat pixFmt;
    int fps;
} VideoEncodeSpec;

static int check_pix_fmt(const AVCodec *codec, enum AVPixelFormat pixFormat) {
    const enum AVPixelFormat *p = codec->pix_fmts;
    while (*p != AV_SAMPLE_FMT_NONE) {
        if (*p == pixFormat) return 1;
        p++;
    }
    return 0;
}


int encode(AVCodecContext *ctx, AVFrame *frame, AVPacket *pkt, FILE *outFile) {
    // 发送数据到编码器
    int ret = avcodec_send_frame(ctx, frame);
    if (ret < 0) {
        ERROR_BUF(ret);
        av_log(NULL, AV_LOG_INFO, "avcodec_send_frame error %s\n", errbuf);
        return ret;
    }
    while (1) {
        // 从编码器中获取编码后的数据
        ret = avcodec_receive_packet(ctx, pkt);
        // packet中已经没有数据了，需要重新发送数据到编码器
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            return 0;
        } else if (ret < 0) {
            ERROR_BUF(ret);
            av_log(NULL, AV_LOG_INFO, "avcodec_receive_packet error %s\n", errbuf);
            return ret;
        }
        // 写入文件
        fwrite(pkt->data, 1, pkt->size, outFile);
        // 释放资源
        av_packet_unref(pkt);
    }
}

void h264Encode(VideoEncodeSpec *in, const char *outFilename) {
    FILE *inFile = NULL;
    FILE *outFile = NULL;
    // 一帧图片的大小
    int imgSize = av_image_get_buffer_size(in->pixFmt, in->width, in->height, 1);
    // 返回结果
    int ret = 0;
    // 编码器
    const AVCodec *codec = NULL;
    // 编码上下文
    AVCodecContext *ctx = NULL;
    AVFrame *frame = NULL; // 用于存放编码前的数据（yuv）
    AVPacket *pkt = NULL; // 用于存放编码后的数据（h264）

    // 获取编码器
    codec = avcodec_find_encoder_by_name("libx264");
    if (!codec) {
        av_log(NULL, AV_LOG_INFO, "avcodec_find_encoder error\n");
        return;
    }

    // 检查编码格式
    if (!check_pix_fmt(codec, in->pixFmt)) {
        av_log(NULL, AV_LOG_INFO, "Encode does not support sample format %s\n",
               av_get_pix_fmt_name(in->pixFmt));
        return;
    }

    // 创建上下文
    ctx = avcodec_alloc_context3(codec);
    if (!ctx) {
        av_log(NULL, AV_LOG_INFO, "avcodec_alloc_context3 error\n");
        return;
    }

    // 设置参数
    ctx->width = in->width;
    ctx->height = in->height;
    ctx->pix_fmt = in->pixFmt;
    // 帧率（1s显示的帧数是in->fps）
    ctx->time_base = (AVRational){1, in->fps};

    // 打开编码器
    ret = avcodec_open2(ctx, codec, NULL);
    if (ret < 0) {
        ERROR_BUF(ret);
        av_log(NULL, AV_LOG_INFO, "avcodec_open2 error %s\n", errbuf);
        goto end;
    }

    // 创建AVFrame
    frame = av_frame_alloc();
    if (!frame) {
        av_log(NULL, AV_LOG_INFO, "av_frame_alloc error\n");
        goto end;
    }
    frame->width = ctx->width;
    frame->height = ctx->height;
    frame->format = ctx->pix_fmt;
    frame->pts = 0;

    // 利用width、height、format创建缓冲区
    ret = av_image_alloc(frame->data,
                         frame->linesize,
                         in->width,
                         in->height,
                         in->pixFmt,
                         1);
    if (ret < 0) {
        ERROR_BUF(ret);
        av_log(NULL, AV_LOG_INFO, "av_frame_get_buffer error %s\n", errbuf);
        goto end;
    }

    // 创建AVPacket
    pkt = av_packet_alloc();
    if (!pkt) {
        av_log(NULL, AV_LOG_INFO, "av_packet_alloc error \n");
        goto end;
    }

    // 打开文件
    inFile = fopen(in->filename, "rb");
    if (inFile == NULL) {
        av_log(NULL, AV_LOG_INFO, "File open error %s\n", in->filename);
        goto end;
    }
    outFile = fopen(outFilename, "wb");
    if (outFile == NULL) {
        av_log(NULL, AV_LOG_INFO, "File open error %s\n", outFilename);
        goto end;
    }

    // 开始编码
    while ((ret = fread(frame->data[0], 1, imgSize, inFile)) > 0) {
        // 编码
        if (encode(ctx, frame, pkt, outFile) < 0) {
            goto end;
        }
        // 设置帧序号
        frame->pts ++;
    }
    // flush编码器
    encode(ctx, NULL, pkt, outFile);

    end:
    fclose(inFile);
    fclose(outFile);
    if (frame){
        av_freep(&frame->data[0]);
        av_frame_free(&frame);
    }
    av_packet_free(&pkt);
    avcodec_free_context(&ctx);
}

// 总是比ffmpeg转换的多1个字节？？？
int main(int argc, char *argv[]) {
    // 输入参数
    VideoEncodeSpec in;
    in.filename = "C:/Users/xu/Desktop/out.yuv";
    in.pixFmt = AV_PIX_FMT_YUV420P;
    in.width = 1920;
    in.height = 1080;
    in.fps = 30;
    h264Encode(&in, "C:/Users/xu/Desktop/out1.h264");
    return 0;
}